package com.ktjiaoyu.travelaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.ktjiaoyu.travelaiagent.agent.enums.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    // 可用的工具集合
    private final ToolCallback[] availableTools;
    // 保存了工具调用信息的响应
    private ChatResponse toolCallChatResponse;
    // 工具调用管理者
    private final ToolCallingManager toolCallingManager;
    // 禁用内置的工具调用机制，自己维护上下文
    private final ChatOptions chatOptions;

    @Override
    public boolean think() {
        // 如果有下一步提示词，则添加为用户消息
        String nextStepPrompt = getNextStepPrompt();
        if (StrUtil.isNotBlank(nextStepPrompt)) {
            UserMessage userMessage = new UserMessage(nextStepPrompt);
            chatMemory.add(this.chatId, userMessage);
        }
        // 获取所有消息
        List<Message> messageList = chatMemory.get(this.chatId, Integer.MAX_VALUE);
        // 传入上下文和选项作为当前提示词
        Prompt prompt = new Prompt(messageList, chatOptions);
        try {
            // 自行调用大模型获得带工具选项的响应
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)  // 注册可用工具
                    // TokenCounterAdvisor依赖chatId，所以这里需要设置chatId
                    .advisors(spec -> spec
                            .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                    // 设置ToolContext信息
                    .toolContext(Map.of("chatId", chatId))
                    .call()
                    .chatResponse();
            // 保存响应，用于act()进行处理
            this.toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            // 输出提示信息
            String result = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            log.info("{}的思考: {}", getName(), result);
            log.info("{}选择了 {} 个工具来使用，具体工具列表如下：", getName(), toolCallList.size());
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s，参数：%s",
                            toolCall.name(),
                            toolCall.arguments())
                    )
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);
            // 如果没有工具调用，则记录助手消息
            if (toolCallList.isEmpty()) {
                chatMemory.add(this.chatId, assistantMessage);
                return false;
            } else {
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        } catch (Exception e) {
            e.printStackTrace();
            log.error("{}的思考过程遇到了问题: {}", getName(), e.getMessage());
            chatMemory.add(this.chatId,
                    new AssistantMessage("处理时遇到错误: " + e.getMessage()));
            return false;
        }
    }

    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具调用";
        }
        // 获取所有消息信息
        List<Message> messages = chatMemory.get(chatId, Integer.MAX_VALUE);
        // 需要调用工具，通过toolCallingManager进行执行
        Prompt prompt = new Prompt(messages, chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        this.chatMemory.clear(this.chatId);
        this.chatMemory.add(this.chatId, toolExecutionResult.conversationHistory());
        // 获得最后一次的工具调用响应信息
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage)
                CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() +
                        " 完成了它的任务！结果: " + response.responseData())
                .collect(Collectors.joining("\n"));
        // 判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));
        if (terminateToolCalled) {
            // 如果调用了终止工具，则设置智能体状态为FINISHED
            setState(AgentState.FINISHED);
        }
        log.info(results);
        return results;
    }
}
